# -*- coding: utf-8 -*-
"""
   File Name：     faceRecognition.py
   Description :  wrapping for face recognition class, instantiation occurs in modules/face_server/utils/api.py
   Author :       KangJiaHui
   date：         2020/01/14
"""

import dlib
import cv2
import numpy as np
import os
import urllib.request
from modules.utils.error import RegisterFaceNum, FaceArea
from modules.utils.utils import image_to_base64, base64_to_image


def calculate_distance(vector1, vector2):
    """
    Calculates Euclidean distance between two vectors.
    :param vector1:vector presents a face feature
    :param vector2:vector presents a face feature
    :return:disance:the Euclidean distance between vector1 and vector2.
    """
    temp = vector1 - vector2
    distance = np.linalg.norm(temp)
    return distance


def download_from_url(filepath, save_dir):
    """
    download file from URL
    :param filepath: str, URL
    :param save_dir: str, save path without filename
    :return: None
    """
    print('\nDownloading file from {}'.format(filepath))
    filename = filepath.split('/')[-1]
    save_path = os.path.join(save_dir, filename)
    urllib.request.urlretrieve(filepath, save_path)
    print('\nSuccessfully downloaded to {}'.format(save_path))


class MatchException(Exception):
    def __init__(self):
        self.msg = "Invalid image input!"

    def __str__(self):
        return str(self.msg)


class FaceRecognition(object):
    def __init__(self):
        self.predictor_path = os.path.join(os.getcwd(), 'modules/face_server/params',
                                           'shape_predictor_68_face_landmarks.dat')
        self.face_rec_model_path = os.path.join(os.getcwd(), 'modules/face_server/params',
                                                'dlib_face_recognition_resnet_model_v1.dat')
        if not os.path.exists('modules/face_server/params'):
            os.makedirs('modules/face_server/params')
        if not os.path.exists(self.predictor_path):
            download_from_url("http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2",
                              os.path.join(os.getcwd(), 'modules/face_server/params'))
            os.system('bzip2 -d modules/face_server/params/shape_predictor_68_face_landmarks.dat.bz2')
        if not os.path.exists(self.face_rec_model_path):
            download_from_url("http://dlib.net/files/dlib_face_recognition_resnet_model_v1.dat.bz2",
                              os.path.join(os.getcwd(), 'modules/face_server/params'))
            os.system('bzip2 -d modules/face_server/params/dlib_face_recognition_resnet_model_v1.dat.bz2')
        self.detector = dlib.get_frontal_face_detector()
        self.sp = dlib.shape_predictor(self.predictor_path)
        self.facerec = dlib.face_recognition_model_v1(self.face_rec_model_path)

    def face_register(self, image, score_reg):
        """
        Registers only one face in one picture.
        :param score_reg: float, the score of the face to be registered should be larger than score_reg
        :param image: array image
        :return: list: the feature vector of the face in input image
        """
        # image = base64_to_image(image_base64)
        # print(image.shape[0], image.shape[1])
        if image.shape[0] > 1000 or image.shape[1] > 1000:
            if image.shape[0] > image.shape[1]:
                proportion = 1000 / image.shape[0]
            else:
                proportion = 1000 / image.shape[1]
            image = cv2.resize(image, None, fx=proportion, fy=proportion, interpolation=cv2.INTER_AREA)
            # print(image.shape[0], image.shape[1])
        faces, scores, idx = self.detector.run(image, 1, score_reg)
        if len(faces) != 1:
            print("there has {} faces!".format(len(faces)))
            raise RegisterFaceNum
        w = faces[0].right() - faces[0].left()
        h = faces[0].bottom() - faces[0].top()
        # print(w/image.shape[1], h/image.shape[0])
        # if w/image.shape[1] + h/image.shape[0] < 0.5:
        #     raise FaceArea
        shape = self.sp(image, faces[0])
        face_chip = dlib.get_face_chip(image, shape)
        feature_vector = list(self.facerec.compute_face_descriptor(face_chip))
        return feature_vector

    def match_identity(self, feature_vector, thresh, score_rec, image_base64=None, path=None):
        """
        Match the input feature vector and vector for each face in one image.
        Once a face matched, it will return True.
        :param score_rec: float, the score of detected faces should be larger than score_rec
        :param feature_vector: list, the feature vector to be matched
        :param thresh: distance between face and matched face should be smaller than thresh
        :param image_base64: image encoded in base64
        :param path: str, indicates an image
        :return:bool: True means face matched, False means face not matched.
        """
        if path:
            image = cv2.imread(path)  # if path and image coexist, then path will cover image
        elif image_base64:
            image = base64_to_image(image_base64)
        else:
            raise MatchException
        if image.shape[0] > 1000 or image.shape[1] > 1000:
            if image.shape[0] > image.shape[1]:
                proportion = 1000 / image.shape[0]
            else:
                proportion = 1000 / image.shape[1]
            image = cv2.resize(image, None, fx=proportion, fy=proportion, interpolation=cv2.INTER_AREA)
        faces, scores, idx = self.detector.run(image, 1, score_rec)
        distance = 1
        for face in faces:
            shape = self.sp(image, face)
            face_chip = dlib.get_face_chip(image, shape)
            face_descriptor = np.array(self.facerec.compute_face_descriptor(face_chip))
            distance = calculate_distance(face_descriptor, np.array(feature_vector))
            if distance < thresh:
                return True, distance  # True means face matched
        return False, distance  # False means face not matched

    @staticmethod
    def match_vectors(feature_vector, candidate_dict, thresh_req, thresh_high=0.3):
        """
        Match the input feature vector and vector for each face in one image.
        Once a face matched, it will return True.
        :param thresh_high: a very strict thresh.
        :param thresh_req: the thresh used in requirement.
        :param candidate_dict:e.x.
            {101:[-0.06939336657524109, 0.11498239636421204, …], 102:[ 0.052934370934963226, 0.004473036155104637, …]}
        :param feature_vector: list, the feature vector to be matched
        :return:bool: True means face matched, False means face not matched.
        """
        result = {}
        feature_vector = np.array(feature_vector)
        for k, v in candidate_dict.items():
            distance = calculate_distance(feature_vector, np.array(v))
            if distance < thresh_high:
                return {k: distance}
            elif thresh_high < distance < thresh_req:
                result[k] = distance
            else:
                continue
        return result

    @staticmethod
    def match_vectors_all(feature_vector, candidate_dict):
        """
        Match the input feature vector and vector for each face in one image.
        Once a face matched, it will return True.
        :param candidate_dict:e.x.
            {101:[-0.06939336657524109, 0.11498239636421204, …], 102:[ 0.052934370934963226, 0.004473036155104637, …]}
        :param feature_vector: list, the feature vector to be matched
        :return:bool: True means face matched, False means face not matched.
        """
        result = {}
        feature_vector = np.array(feature_vector)
        for k, v in candidate_dict.items():
            distance = calculate_distance(feature_vector, np.array(v))
            result[k] = distance
        return result

